// This code is part of Qiskit.
//
// (C) Copyright IBM 2024
//
// This code is licensed under the Apache License, Version 2.0. You may
// obtain a copy of this license in the LICENSE.txt file in the root directory
// of this source tree or at http://www.apache.org/licenses/LICENSE-2.0.
//
// Any modifications or derivative works of this code must retain this
// copyright notice, and modified files need to carry a notice indicating
// that they have been altered from the originals.

use hashbrown::{HashMap, HashSet};
use ndarray::linalg::kron;
use ndarray::Array2;
use num_complex::Complex64;
use num_complex::ComplexFloat;
use qiskit_circuit::object_registry::PyObjectAsKey;
use smallvec::SmallVec;
use std::fmt::Debug;

use numpy::PyReadonlyArray2;
use pyo3::exceptions::PyRuntimeError;
use pyo3::intern;
use pyo3::prelude::*;
use pyo3::types::{PyBool, PyDict, PyTuple};
use pyo3::BoundObject;

use qiskit_circuit::circuit_instruction::OperationFromPython;
use qiskit_circuit::dag_node::DAGOpNode;
use qiskit_circuit::imports::QI_OPERATOR;
use qiskit_circuit::object_registry::ObjectRegistry;
use qiskit_circuit::operations::{
    Operation, OperationRef, Param, StandardGate, STANDARD_GATE_SIZE,
};
use qiskit_circuit::{Clbit, Qubit};
use qiskit_quantum_info::unitary_compose;
use thiserror::Error;

use crate::gate_metrics;
use crate::standard_gates_commutations;
use crate::QiskitError;

#[derive(Error, Debug)]
pub enum CommutationError {
    #[error("Einsum error: {0}")]
    EinsumError(&'static str),
    #[error("First instruction must have at most as many qubits as the second instruction")]
    FirstInstructionTooLarge,
    #[error("Invalid hash value: NaN or inf")]
    HashingNaN,
    #[error("Invalid hash type: parameterized")]
    HashingParameter,
}

impl From<CommutationError> for PyErr {
    fn from(value: CommutationError) -> Self {
        match value {
            // For backward compatibility we keep these two errors as QiskitErrors
            CommutationError::HashingParameter | CommutationError::FirstInstructionTooLarge => {
                QiskitError::new_err(value.to_string())
            }
            _ => PyRuntimeError::new_err(value.to_string()),
        }
    }
}

const fn build_supported_ops() -> [bool; STANDARD_GATE_SIZE] {
    let mut lut = [false; STANDARD_GATE_SIZE];
    lut[StandardGate::RXX as usize] = true;
    lut[StandardGate::RYY as usize] = true;
    lut[StandardGate::RZZ as usize] = true;
    lut[StandardGate::RZX as usize] = true;
    lut[StandardGate::H as usize] = true;
    lut[StandardGate::X as usize] = true;
    lut[StandardGate::Y as usize] = true;
    lut[StandardGate::Z as usize] = true;
    lut[StandardGate::SX as usize] = true;
    lut[StandardGate::SXdg as usize] = true;
    lut[StandardGate::T as usize] = true;
    lut[StandardGate::Tdg as usize] = true;
    lut[StandardGate::S as usize] = true;
    lut[StandardGate::Sdg as usize] = true;
    lut[StandardGate::CX as usize] = true;
    lut[StandardGate::CY as usize] = true;
    lut[StandardGate::CZ as usize] = true;
    lut[StandardGate::Swap as usize] = true;
    lut[StandardGate::ISwap as usize] = true;
    lut[StandardGate::ECR as usize] = true;
    lut[StandardGate::CCX as usize] = true;
    lut[StandardGate::CSwap as usize] = true;
    lut
}

static SUPPORTED_OP: [bool; STANDARD_GATE_SIZE] = build_supported_ops();

// Map rotation gates to their generators (or to ``None`` if we cannot currently efficiently
// represent the generator in Rust and store the commutation relation in the commutation dictionary)
// and their pi-periodicity. Here we mean a gate is n-pi periodic, if for angles that are
// multiples of n*pi, the gate is equal to the identity up to a global phase.
// E.g. RX is generated by X and 2-pi periodic, while CRX is generated by CX and 4-pi periodic.
const fn build_supported_rotations() -> [Option<Option<StandardGate>>; STANDARD_GATE_SIZE] {
    let mut lut = [None; STANDARD_GATE_SIZE];
    lut[StandardGate::RX as usize] = Some(Some(StandardGate::X));
    lut[StandardGate::RY as usize] = Some(Some(StandardGate::Y));
    lut[StandardGate::RZ as usize] = Some(Some(StandardGate::Z));
    lut[StandardGate::Phase as usize] = Some(Some(StandardGate::Z));
    lut[StandardGate::U1 as usize] = Some(Some(StandardGate::Z));
    lut[StandardGate::CRX as usize] = Some(Some(StandardGate::CX));
    lut[StandardGate::CRY as usize] = Some(Some(StandardGate::CY));
    lut[StandardGate::CRZ as usize] = Some(Some(StandardGate::CZ));
    lut[StandardGate::CPhase as usize] = Some(Some(StandardGate::CZ));
    // RXXGate, RYYGate, RZXGate, and RZZGate are supported by the commutation dictionary
    lut[StandardGate::RXX as usize] = Some(None);
    lut[StandardGate::RYY as usize] = Some(None);
    lut[StandardGate::RZX as usize] = Some(None);
    lut[StandardGate::RZZ as usize] = Some(None);
    lut
}

static SUPPORTED_ROTATIONS: [Option<Option<StandardGate>>; STANDARD_GATE_SIZE] =
    build_supported_rotations();

fn get_bits_from_py<T>(
    py_bits1: &Bound<'_, PyTuple>,
    py_bits2: &Bound<'_, PyTuple>,
) -> PyResult<(Vec<T>, Vec<T>)>
where
    T: From<u32> + Copy,
    u32: From<T>,
{
    // Using `PyObjectAsKey` here is a total hack, but this is a short-term workaround before a
    // larger refactor of the commutation checker.
    let mut registry: ObjectRegistry<T, PyObjectAsKey> = ObjectRegistry::new();

    for bit in py_bits1.iter().chain(py_bits2.iter()) {
        registry.add(bit.into(), false)?;
    }

    Ok((
        registry
            .map_objects(py_bits1.iter().map(|bit| bit.into()))?
            .collect(),
        registry
            .map_objects(py_bits2.iter().map(|bit| bit.into()))?
            .collect(),
    ))
}

/// This is the internal structure for the Python CommutationChecker class
/// It handles the actual commutation checking, cache management, and library
/// lookups. It's not meant to be a public facing Python object though and only used
/// internally by the Python class.
#[pyclass(module = "qiskit._accelerate.commutation_checker")]
pub struct CommutationChecker {
    library: CommutationLibrary,
    cache_max_entries: usize,
    cache: HashMap<(String, String), CommutationCacheEntry>,
    current_cache_entries: usize,
    #[pyo3(get)]
    gates: Option<HashSet<String>>,
}

#[pymethods]
impl CommutationChecker {
    #[pyo3(signature = (standard_gate_commutations=None, cache_max_entries=1_000_000, gates=None))]
    #[new]
    fn py_new(
        standard_gate_commutations: Option<Bound<PyAny>>,
        cache_max_entries: usize,
        gates: Option<HashSet<String>>,
    ) -> Self {
        let library = CommutationLibrary::py_new(standard_gate_commutations);
        CommutationChecker::new(Some(library), cache_max_entries, gates)
    }

    #[pyo3(signature=(op1, op2, max_num_qubits=3, approximation_degree=1.))]
    fn commute_nodes(
        &mut self,
        py: Python,
        op1: &DAGOpNode,
        op2: &DAGOpNode,
        max_num_qubits: u32,
        approximation_degree: f64,
    ) -> PyResult<bool> {
        let (qargs1, qargs2) = get_bits_from_py::<Qubit>(
            op1.instruction.qubits.bind(py),
            op2.instruction.qubits.bind(py),
        )?;
        let (cargs1, cargs2) = get_bits_from_py::<Clbit>(
            op1.instruction.clbits.bind(py),
            op2.instruction.clbits.bind(py),
        )?;

        Ok(self.commute(
            &op1.instruction.operation.view(),
            &op1.instruction.params,
            &qargs1,
            &cargs1,
            &op2.instruction.operation.view(),
            &op2.instruction.params,
            &qargs2,
            &cargs2,
            max_num_qubits,
            approximation_degree,
        )?)
    }

    #[pyo3(name="commute", signature=(op1, qargs1, cargs1, op2, qargs2, cargs2, max_num_qubits=3, approximation_degree=1.))]
    #[allow(clippy::too_many_arguments)]
    fn py_commute(
        &mut self,
        op1: OperationFromPython,
        qargs1: &Bound<'_, PyTuple>,
        cargs1: &Bound<'_, PyTuple>,
        op2: OperationFromPython,
        qargs2: &Bound<'_, PyTuple>,
        cargs2: &Bound<'_, PyTuple>,
        max_num_qubits: u32,
        approximation_degree: f64,
    ) -> PyResult<bool> {
        let (qargs1, qargs2) = get_bits_from_py::<Qubit>(qargs1, qargs2)?;
        let (cargs1, cargs2) = get_bits_from_py::<Clbit>(cargs1, cargs2)?;

        Ok(self.commute(
            &op1.operation.view(),
            &op1.params,
            &qargs1,
            &cargs1,
            &op2.operation.view(),
            &op2.params,
            &qargs2,
            &cargs2,
            max_num_qubits,
            approximation_degree,
        )?)
    }

    /// Return the current number of cache entries
    fn num_cached_entries(&self) -> usize {
        self.current_cache_entries
    }

    /// Clear the cache
    fn clear_cached_commutations(&mut self) {
        self.clear_cache()
    }

    fn __getstate__(&self, py: Python) -> PyResult<Py<PyDict>> {
        let out_dict = PyDict::new(py);
        out_dict.set_item("cache_max_entries", self.cache_max_entries)?;
        out_dict.set_item("current_cache_entries", self.current_cache_entries)?;
        let cache_dict = PyDict::new(py);
        for (key, value) in &self.cache {
            cache_dict.set_item(key, commutation_entry_to_pydict(py, value)?)?;
        }
        out_dict.set_item("cache", cache_dict)?;
        out_dict.set_item("library", self.library.library.clone().into_pyobject(py)?)?;
        out_dict.set_item("gates", self.gates.clone())?;
        Ok(out_dict.unbind())
    }

    fn __setstate__(&mut self, py: Python, state: Py<PyAny>) -> PyResult<()> {
        let dict_state = state.downcast_bound::<PyDict>(py)?;
        self.cache_max_entries = dict_state
            .get_item("cache_max_entries")?
            .unwrap()
            .extract()?;
        self.current_cache_entries = dict_state
            .get_item("current_cache_entries")?
            .unwrap()
            .extract()?;
        self.library = CommutationLibrary {
            library: dict_state.get_item("library")?.unwrap().extract()?,
        };
        let raw_cache: Bound<PyDict> = dict_state.get_item("cache")?.unwrap().extract()?;
        self.cache = HashMap::with_capacity(raw_cache.len());
        for (key, value) in raw_cache.iter() {
            let value_dict: &Bound<PyDict> = value.downcast()?;
            self.cache.insert(
                key.extract()?,
                commutation_cache_entry_from_pydict(value_dict)?,
            );
        }
        self.gates = dict_state.get_item("gates")?.unwrap().extract()?;
        Ok(())
    }
}

impl CommutationChecker {
    /// Create a new commutation checker.
    ///
    /// # Arguments
    ///
    /// - `library`: An optional existing [CommutationLibrary] with cached entries.
    /// - `cache_max_entries`: The maximum size of the cache.
    /// - `gates`: An optional set of gates (by name) to check commutations for. If `None`,
    ///   commutation is cached and checked for all gates.
    pub fn new(
        library: Option<CommutationLibrary>,
        cache_max_entries: usize,
        gates: Option<HashSet<String>>,
    ) -> Self {
        // Initialize sets before they are used in the commutation checker
        CommutationChecker {
            library: library.unwrap_or(CommutationLibrary { library: None }),
            cache: HashMap::new(),
            cache_max_entries,
            current_cache_entries: 0,
            gates,
        }
    }

    #[allow(clippy::too_many_arguments)]
    pub fn commute(
        &mut self,
        op1: &OperationRef,
        params1: &[Param],
        qargs1: &[Qubit],
        cargs1: &[Clbit],
        op2: &OperationRef,
        params2: &[Param],
        qargs2: &[Qubit],
        cargs2: &[Clbit],
        max_num_qubits: u32,
        approximation_degree: f64,
    ) -> Result<bool, CommutationError> {
        // If the average gate infidelity is below this tolerance, they commute. The tolerance
        // is set to max(1e-12, 1 - approximation_degree), to account for roundoffs and for
        // consistency with other places in Qiskit.
        let tol = 1e-12_f64.max(1. - approximation_degree);

        // if we have rotation gates, we attempt to map them to their generators, for example
        // RX -> X or CPhase -> CZ
        let (op1_gate, params1, trivial1) = map_rotation(op1, params1, tol);
        if trivial1 {
            return Ok(true);
        }
        let op1 = if let Some(gate) = op1_gate {
            &OperationRef::StandardGate(gate)
        } else {
            op1
        };
        let (op2_gate, params2, trivial2) = map_rotation(op2, params2, tol);
        if trivial2 {
            return Ok(true);
        }
        let op2 = if let Some(gate) = op2_gate {
            &OperationRef::StandardGate(gate)
        } else {
            op2
        };

        if let Some(gates) = &self.gates {
            if !gates.is_empty() && (!gates.contains(op1.name()) || !gates.contains(op2.name())) {
                return Ok(false);
            }
        }

        let commutation: Option<bool> = commutation_precheck(
            op1,
            params1,
            qargs1,
            cargs1,
            op2,
            params2,
            qargs2,
            cargs2,
            max_num_qubits,
        );
        if let Some(is_commuting) = commutation {
            return Ok(is_commuting);
        }

        let reversed = if op1.num_qubits() != op2.num_qubits() {
            op1.num_qubits() > op2.num_qubits()
        } else {
            (op1.name().len(), op1.name()) >= (op2.name().len(), op2.name())
        };
        let (first_params, second_params) = if reversed {
            (params2, params1)
        } else {
            (params1, params2)
        };
        let (first_op, second_op) = if reversed { (op2, op1) } else { (op1, op2) };
        let (first_qargs, second_qargs) = if reversed {
            (qargs2, qargs1)
        } else {
            (qargs1, qargs2)
        };

        // For our cache to work correctly, we require the gate's definition to only depend on the
        // ``params`` attribute. This cannot be guaranteed for custom gates, so we only check
        // the cache for
        //  * gates we know are in the cache (SUPPORTED_OPS), or
        //  * standard gates with float params (otherwise we cannot cache them)
        let is_cachable = |op: &OperationRef, params: &[Param]| {
            if let OperationRef::StandardGate(gate) = op {
                SUPPORTED_OP[(*gate) as usize]
                    || params.iter().all(|p| matches!(p, Param::Float(_)))
            } else {
                false
            }
        };
        let check_cache =
            is_cachable(first_op, first_params) && is_cachable(second_op, second_params);

        if !check_cache {
            return self.commute_matmul(
                first_op,
                first_params,
                first_qargs,
                second_op,
                second_params,
                second_qargs,
                tol,
            );
        }

        // Query commutation library
        let relative_placement = get_relative_placement(first_qargs, second_qargs);
        if let Some(is_commuting) =
            self.library
                .check_commutation_entries(first_op, second_op, &relative_placement)
        {
            return Ok(is_commuting);
        }

        // Query cache
        let key1 = hashable_params(first_params)?;
        let key2 = hashable_params(second_params)?;
        if let Some(commutation_dict) = self
            .cache
            .get(&(first_op.name().to_string(), second_op.name().to_string()))
        {
            let hashes = (key1.clone(), key2.clone());
            if let Some(commutation) = commutation_dict.get(&(relative_placement.clone(), hashes)) {
                return Ok(*commutation);
            }
        }

        // Perform matrix multiplication to determine commutation
        let is_commuting = self.commute_matmul(
            first_op,
            first_params,
            first_qargs,
            second_op,
            second_params,
            second_qargs,
            tol,
        )?;

        // TODO: implement a LRU cache for this
        if self.current_cache_entries >= self.cache_max_entries {
            self.clear_cache();
        }
        // Cache results from is_commuting
        self.cache
            .entry((first_op.name().to_string(), second_op.name().to_string()))
            .and_modify(|entries| {
                let key = (relative_placement.clone(), (key1.clone(), key2.clone()));
                entries.insert(key, is_commuting);
                self.current_cache_entries += 1;
            })
            .or_insert_with(|| {
                let mut entries = HashMap::with_capacity(1);
                let key = (relative_placement, (key1, key2));
                entries.insert(key, is_commuting);
                self.current_cache_entries += 1;
                entries
            });
        Ok(is_commuting)
    }

    #[allow(clippy::too_many_arguments)]
    fn commute_matmul(
        &self,
        first_op: &OperationRef,
        first_params: &[Param],
        first_qargs: &[Qubit],
        second_op: &OperationRef,
        second_params: &[Param],
        second_qargs: &[Qubit],
        tol: f64,
    ) -> Result<bool, CommutationError> {
        // Compute relative positioning of qargs of the second gate to the first gate.
        // Since the qargs come out the same BitData, we already know there are no accidential
        // bit-duplications, but this code additionally maps first_qargs to [0..n] and then
        // computes second_qargs relative to that. For example, it performs the mappings
        //  (first_qargs, second_qargs) = ( [1, 2], [0, 2] ) --> ( [0, 1], [2, 1] )
        //  (first_qargs, second_qargs) = ( [1, 2, 0], [0, 3, 4] ) --> ( [0, 1, 2], [2, 3, 4] )
        // This re-shuffling is done to compute the correct kronecker product later.
        let mut qarg: HashMap<&Qubit, Qubit> = HashMap::from_iter(
            first_qargs
                .iter()
                .enumerate()
                .map(|(i, q)| (q, Qubit::new(i))),
        );
        let mut num_qubits = first_qargs.len() as u32;
        for q in second_qargs {
            if !qarg.contains_key(q) {
                qarg.insert(q, Qubit(num_qubits));
                num_qubits += 1;
            }
        }

        let first_qarg: Vec<Qubit> = Vec::from_iter((0..first_qargs.len() as u32).map(Qubit));
        let second_qarg: Vec<Qubit> = second_qargs.iter().map(|q| qarg[q]).collect();

        if first_qarg.len() > second_qarg.len() {
            return Err(CommutationError::FirstInstructionTooLarge);
        };
        let first_mat = match get_matrix(first_op, first_params) {
            Some(matrix) => matrix,
            None => return Ok(false),
        };

        let second_mat = match get_matrix(second_op, second_params) {
            Some(matrix) => matrix,
            None => return Ok(false),
        };

        // TODO Optimize this bit to avoid unnecessary Kronecker products:
        //  1. We currently sort the operations for the cache by operation size, putting the
        //     *smaller* operation first: (smaller op, larger op)
        //  2. This code here expands the first op to match the second -- hence we always
        //     match the operator sizes.
        // This whole extension logic could be avoided since we know the second one is larger.
        let extra_qarg2 = num_qubits - first_qarg.len() as u32;
        let first_mat = if extra_qarg2 > 0 {
            let id_op = Array2::<Complex64>::eye(usize::pow(2, extra_qarg2));
            kron(&id_op, &first_mat)
        } else {
            first_mat
        };

        // the 1 qubit case cannot happen, since that would already have been captured
        // by the previous if clause; first_qarg == second_qarg (if they overlap they must
        // be the same)
        let op12 = match unitary_compose::compose(
            &first_mat.view(),
            &second_mat.view(),
            &second_qarg,
            false,
        ) {
            Ok(matrix) => matrix,
            Err(e) => return Err(CommutationError::EinsumError(e)),
        };
        let op21 = match unitary_compose::compose(
            &first_mat.view(),
            &second_mat.view(),
            &second_qarg,
            true,
        ) {
            Ok(matrix) => matrix,
            Err(e) => return Err(CommutationError::EinsumError(e)),
        };
        let (fid, phase) = gate_metrics::gate_fidelity(&op12.view(), &op21.view(), None);

        // we consider the gates as commuting if the process fidelity of
        // AB (BA)^\dagger is approximately the identity and there is no global phase difference
        // let dim = op12.ncols() as f64;
        // let matrix_tol = tol * dim.powi(2);
        let matrix_tol = tol;
        Ok(phase.abs() <= tol && (1.0 - fid).abs() <= matrix_tol)
    }

    fn clear_cache(&mut self) {
        self.cache.clear();
        self.current_cache_entries = 0;
    }
}

#[allow(clippy::too_many_arguments)]
fn commutation_precheck(
    op1: &OperationRef,
    params1: &[Param],
    qargs1: &[Qubit],
    cargs1: &[Clbit],
    op2: &OperationRef,
    params2: &[Param],
    qargs2: &[Qubit],
    cargs2: &[Clbit],
    max_num_qubits: u32,
) -> Option<bool> {
    if op1.control_flow() || op2.control_flow() {
        return Some(false);
    }

    // assuming the number of involved qubits to be small, this might be faster than set operations
    if !qargs1.iter().any(|e| qargs2.contains(e)) && !cargs1.iter().any(|e| cargs2.contains(e)) {
        return Some(true);
    }

    if qargs1.len() > max_num_qubits as usize || qargs2.len() > max_num_qubits as usize {
        return Some(false);
    }

    if let OperationRef::StandardGate(gate_1) = op1 {
        if let OperationRef::StandardGate(gate_2) = op2 {
            if SUPPORTED_OP[(*gate_1) as usize] && SUPPORTED_OP[(*gate_2) as usize] {
                return None;
            }
        }
    }

    if matches!(
        op1,
        OperationRef::StandardInstruction(_) | OperationRef::Instruction(_)
    ) || matches!(
        op2,
        OperationRef::StandardInstruction(_) | OperationRef::Instruction(_)
    ) {
        return Some(false);
    }

    if is_parameterized(params1) || is_parameterized(params2) {
        return Some(false);
    }

    None
}

fn get_matrix(operation: &OperationRef, params: &[Param]) -> Option<Array2<Complex64>> {
    if let Some(matrix) = operation.matrix(params) {
        Some(matrix)
    } else {
        match operation {
            OperationRef::Gate(gate) => Python::attach(|py| -> Option<_> {
                Some(
                    QI_OPERATOR
                        .get_bound(py)
                        .call1((gate.gate.clone_ref(py),))
                        .ok()?
                        .getattr(intern!(py, "data"))
                        .ok()?
                        .extract::<PyReadonlyArray2<Complex64>>()
                        .ok()?
                        .as_array()
                        .to_owned(),
                )
            }),
            OperationRef::Operation(operation) => Python::attach(|py| -> Option<_> {
                Some(
                    QI_OPERATOR
                        .get_bound(py)
                        .call1((operation.operation.clone_ref(py),))
                        .ok()?
                        .getattr(intern!(py, "data"))
                        .ok()?
                        .extract::<PyReadonlyArray2<Complex64>>()
                        .ok()?
                        .as_array()
                        .to_owned(),
                )
            }),
            _ => None,
        }
    }
}

fn is_parameterized(params: &[Param]) -> bool {
    params
        .iter()
        .any(|x| matches!(x, Param::ParameterExpression(_)))
}

/// Check if a given operation can be mapped onto a generator.
///
/// If ``op`` is in the ``SUPPORTED_ROTATIONS`` hashmap, it is a rotation and we
///   (1) check whether the rotation is so small (modulo pi) that we assume it is the
///       identity and it commutes trivially with every other operation
///   (2) otherwise, we check whether a generator of the rotation is given (e.g. X for RX)
///       and we return the generator
///
/// Returns (operation, parameters, commutes_trivially).
fn map_rotation<'a>(
    op: &'a OperationRef<'a>,
    params: &'a [Param],
    tol: f64,
) -> (Option<StandardGate>, &'a [Param], bool) {
    if let OperationRef::StandardGate(gate) = op {
        if let Some(generator) = SUPPORTED_ROTATIONS[(*gate) as usize] {
            // If the rotation angle is below the tolerance, the gate is assumed to
            // commute with everything, and we simply return the operation with the flag that
            // it commutes trivially.
            if let Param::Float(angle) = params[0] {
                let (tr_over_dim, dim) = gate_metrics::rotation_trace_and_dim(*gate, angle)
                    .expect("All rotation should be covered at this point");
                let gate_fidelity = tr_over_dim.abs().powi(2);
                let process_fidelity = (dim * gate_fidelity + 1.) / (dim + 1.);
                if (1. - process_fidelity).abs() <= tol {
                    return (Some(*gate), params, true);
                };
            };

            // Otherwise we need to cover two cases -- either a generator is given, in which case
            // we return it, or we don't have a generator yet, but we know we have the operation
            // stored in the commutation library. For example, RXX does not have a generator in Rust
            // yet (PauliGate is not in Rust currently), but it is stored in the library, so we
            // can strip the parameters and just return the gate.
            if let Some(gate) = generator {
                return (Some(gate), &[], false);
            };
        }
    }
    (None, params, false)
}

fn get_relative_placement(
    first_qargs: &[Qubit],
    second_qargs: &[Qubit],
) -> SmallVec<[Option<Qubit>; 2]> {
    let mut qubits_g2: HashMap<&Qubit, Qubit> = HashMap::with_capacity(second_qargs.len());
    second_qargs
        .iter()
        .enumerate()
        .for_each(|(i_g1, q_g1)| unsafe {
            qubits_g2.insert_unique_unchecked(q_g1, Qubit::new(i_g1));
        });

    first_qargs
        .iter()
        .map(|q_g0| qubits_g2.get(q_g0).copied())
        .collect()
}

#[derive(Clone, Debug)]
#[pyclass]
pub struct CommutationLibrary {
    pub library: Option<HashMap<(String, String), CommutationLibraryEntry>>,
}

impl CommutationLibrary {
    pub(crate) fn with_capacity(size: usize) -> Self {
        CommutationLibrary {
            library: Some(HashMap::with_capacity(size)),
        }
    }

    pub(crate) fn add_entry(&mut self, key: [&str; 2], value: CommutationLibraryEntry) {
        if let Some(library) = &mut self.library {
            library.insert((key[0].to_string(), key[1].to_string()), value);
        } else {
            let mut library = HashMap::new();
            library.insert((key[0].to_string(), key[1].to_string()), value);
            self.library = Some(library);
        }
    }

    fn check_commutation_entries(
        &self,
        first_op: &OperationRef,
        second_op: &OperationRef,
        relative_placement: &SmallVec<[Option<Qubit>; 2]>,
    ) -> Option<bool> {
        if let Some(library) = &self.library {
            match library.get(&(first_op.name().to_string(), second_op.name().to_string())) {
                Some(CommutationLibraryEntry::Commutes(b)) => Some(*b),
                Some(CommutationLibraryEntry::QubitMapping(qm)) => {
                    qm.get(relative_placement).copied()
                }
                _ => None,
            }
        } else {
            None
        }
    }
}

#[pymethods]
impl CommutationLibrary {
    #[new]
    #[pyo3(signature=(py_any=None))]
    fn py_new(py_any: Option<Bound<PyAny>>) -> Self {
        match py_any {
            Some(pyob) => CommutationLibrary {
                library: pyob
                    .extract::<HashMap<(String, String), CommutationLibraryEntry>>()
                    .ok(),
            },
            None => CommutationLibrary {
                library: Some(HashMap::new()),
            },
        }
    }
}

#[derive(Clone, Debug)]
pub enum CommutationLibraryEntry {
    Commutes(bool),
    QubitMapping(HashMap<SmallVec<[Option<Qubit>; 2]>, bool>),
}

impl<'py> IntoPyObject<'py> for CommutationLibraryEntry {
    type Target = PyAny;
    type Output = Bound<'py, Self::Target>;
    type Error = PyErr;

    fn into_pyobject(self, py: Python<'py>) -> Result<Self::Output, Self::Error> {
        let py_out = match self {
            CommutationLibraryEntry::Commutes(b) => b.into_pyobject(py)?.into_bound().into_any(),
            CommutationLibraryEntry::QubitMapping(qm) => {
                let out = PyDict::new(py);
                for (k, v) in qm {
                    let key = PyTuple::new(py, k.iter().map(|q| q.map(|t| t.0)))?;
                    let value = PyBool::new(py, v);
                    out.set_item(key, value)?;
                }
                out.into_any()
            }
        };
        Ok(py_out)
    }
}

impl<'py> FromPyObject<'py> for CommutationLibraryEntry {
    fn extract_bound(b: &Bound<'py, PyAny>) -> Result<Self, PyErr> {
        if let Ok(b) = b.extract::<bool>() {
            return Ok(CommutationLibraryEntry::Commutes(b));
        }
        let dict = b.downcast::<PyDict>()?;
        let mut ret = hashbrown::HashMap::with_capacity(dict.len());
        for (k, v) in dict {
            let raw_key: SmallVec<[Option<u32>; 2]> = k.extract()?;
            let v: bool = v.extract()?;
            let key = raw_key.into_iter().map(|key| key.map(Qubit)).collect();
            ret.insert(key, v);
        }
        Ok(CommutationLibraryEntry::QubitMapping(ret))
    }
}

type CacheKey = (
    SmallVec<[Option<Qubit>; 2]>,
    (SmallVec<[ParameterKey; 3]>, SmallVec<[ParameterKey; 3]>),
);

type CommutationCacheEntry = HashMap<CacheKey, bool>;

fn commutation_entry_to_pydict(py: Python, entry: &CommutationCacheEntry) -> PyResult<Py<PyDict>> {
    let out_dict = PyDict::new(py);
    for (k, v) in entry.iter() {
        let qubits = PyTuple::new(py, k.0.iter().map(|q| q.map(|t| t.0)))?;
        let params0 = PyTuple::new(py, k.1 .0.iter().map(|pk| pk.0))?;
        let params1 = PyTuple::new(py, k.1 .1.iter().map(|pk| pk.0))?;
        out_dict.set_item(
            PyTuple::new(py, [qubits, PyTuple::new(py, [params0, params1])?])?,
            PyBool::new(py, *v),
        )?;
    }
    Ok(out_dict.unbind())
}

fn commutation_cache_entry_from_pydict(dict: &Bound<PyDict>) -> PyResult<CommutationCacheEntry> {
    let mut ret = hashbrown::HashMap::with_capacity(dict.len());
    for (k, v) in dict {
        let raw_key: CacheKeyRaw = k.extract()?;
        let qubits = raw_key.0.iter().map(|q| q.map(Qubit)).collect();
        let params0: SmallVec<_> = raw_key.1 .0;
        let params1: SmallVec<_> = raw_key.1 .1;
        let v: bool = v.extract()?;
        ret.insert((qubits, (params0, params1)), v);
    }
    Ok(ret)
}

type CacheKeyRaw = (
    SmallVec<[Option<u32>; 2]>,
    (SmallVec<[ParameterKey; 3]>, SmallVec<[ParameterKey; 3]>),
);

/// This newtype wraps a f64 to make it hashable so we can cache parameterized gates
/// based on the parameter value (assuming it's a float angle). However, Rust doesn't do
/// this by default and there are edge cases to track around it's usage. The biggest one
/// is this does not work with f64::NAN, f64::INFINITY, or f64::NEG_INFINITY
/// If you try to use these values with this type they will not work as expected.
/// This should only be used with the cache hashmap's keys and not used beyond that.
#[derive(Debug, Copy, Clone, PartialEq, FromPyObject)]
struct ParameterKey(f64);

impl ParameterKey {
    fn key(&self) -> u64 {
        // If we get a -0 the to_bits() return is not equivalent to 0
        // because -0 has the sign bit set we'd be hashing 9223372036854775808
        // and be storing it separately from 0. So this normalizes all 0s to
        // be represented by 0
        if self.0 == 0. {
            0
        } else {
            self.0.to_bits()
        }
    }
}

impl std::hash::Hash for ParameterKey {
    fn hash<H>(&self, state: &mut H)
    where
        H: std::hash::Hasher,
    {
        self.key().hash(state)
    }
}

impl Eq for ParameterKey {}

fn hashable_params(params: &[Param]) -> Result<SmallVec<[ParameterKey; 3]>, CommutationError> {
    params
        .iter()
        .map(|x| {
            if let Param::Float(x) = x {
                // NaN and Infinity (negative or positive) are not valid
                // parameter values and our hacks to store parameters in
                // the cache HashMap don't take these into account. So return
                // an error to Python if we encounter these values.
                if x.is_nan() || x.is_infinite() {
                    Err(CommutationError::HashingNaN)
                } else {
                    Ok(ParameterKey(*x))
                }
            } else {
                Err(CommutationError::HashingParameter)
            }
        })
        .collect()
}

#[pyfunction]
pub fn get_standard_commutation_checker() -> CommutationChecker {
    let library = standard_gates_commutations::get_commutation_library();
    CommutationChecker {
        library,
        cache_max_entries: 1_000_000,
        cache: HashMap::new(),
        current_cache_entries: 0,
        gates: None,
    }
}

pub fn commutation_checker(m: &Bound<PyModule>) -> PyResult<()> {
    m.add_class::<CommutationLibrary>()?;
    m.add_class::<CommutationChecker>()?;
    m.add_wrapped(wrap_pyfunction!(get_standard_commutation_checker))?;
    Ok(())
}
